# _*_ coding: utf-8 _*_
import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import sample, randint, shuffle, uniform
from time import time, gmtime, strftime, localtime
import argparse

from matplotlib import font_manager

from VCMTSP import Instance
from matplotlib import rcParams
from matplotlib import font_manager

path = '../LinBiolinum_R.otf'
font_manager.fontManager.addfont(path)
prop = font_manager.FontProperties(fname=path)
plt.rcParams['font.family'] = prop.get_name()

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
# plt.rc('font', family='Linux Biolinum O', weight='normal', size=12)


ACM_Blue = '#0182ac'
ACM_Yellow = '#ffd600'
ACM_Orange = '#fc9200'
ACM_Red = '#e33532'
ACM_Light_Blue = '#90cef2'
ACM_Green = '#adbf00'
ACM_Purple = '#78006b'
ACM_Dark_Blue = '#184c8c'

parser = argparse.ArgumentParser()
parser.add_argument("--case", required=True, type=str, help="The case name!")
parser.add_argument("-m", required=True, type=int, help="The number of salesmen!")
args = parser.parse_args()

m = args.m
case = args.case

instance = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
Distance = instance.distance  # distance matrix
Cost = instance.cost  # traveling cost of a single path
A = instance.A  # accessibility matrix
D = instance.D  # depot
V = deepcopy(instance.cities[1:]).tolist()  # city set, exclude depot 1import numpy as np


def make_dir(path):
    path = path.strip()
    path = path.rstrip("\\")
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        return True
    else:
        print(path + ' already exists.')
        return False


def normalization(x):
    if np.min(x) != np.max(x):
        return (x - np.min(x)) / (np.max(x) - np.min(x))
    else:
        return (x - np.min(x)) + 1


def find_nearest(city, roll):
    tar, min_dis = -1, 1e99
    for c in roll:
        if min_dis > Distance[city][c]:
            min_dis = Distance[city][c]
            tar = c
    assert tar != -1
    return tar, min_dis


def salesmen_of(city):
    return np.where(A[city, :] == 1)[0].tolist()


def cities_of(salesman):
    return np.where(A[:, salesman] == 1)[0].tolist()


def transition(curr, roll):
    """
    @param curr: the current city where the ant is at
    @param roll: the roll list of cities accessible to the ant
    @return: the next city to reach
    """
    if len(roll) == 0:
        return 0
    roll = list(roll)
    prob = np.power(Pheromone[curr][roll], alpha) * np.power(1 / Distance[curr][roll], beta)
    prob = prob / prob.sum()
    if uniform(0, 1) < q_zero:
        city = roll[int(np.argmax(prob))]
    else:
        city = np.random.choice(a=roll, size=1, p=prob)[0]
    return city


def pheromone_diffusion(route, scope, delta):
    # scope = 1
    # plot = plt
    # dot_size = round(np.log2(instance.n)) * 2
    # x, y = instance.x, instance.y
    # plot.scatter(x[1:], y[1:], s=dot_size, c='grey')
    # plot.xticks([])
    # plot.yticks([])
    # plot.plot(x[route], y[route], c='k', alpha=0.6, lw=3, label='Source trail')

    Dis = deepcopy(Distance)
    for idx in range(1, len(route) - 1):
        a, b = route[idx-1], route[idx]
        Pheromone[a][b] = (1 - rho) * Pheromone[a][b] + rho * delta
        Pheromone[b][a] = (1 - rho) * Pheromone[b][a] + rho * delta
        c = route[idx+1]
        Dis[b][a] = 1e99
        Dis[b][b] = 1e99
        Dis[b][c] = 1e99
        neighbors = np.argsort(Dis[b])[1:scope+1]
        diffusion = delta * 1
        # update pheromone on neighbor edges
        for n in neighbors:
            Pheromone[n][b] = (1 - rho) * Pheromone[a][b] + rho * diffusion
            Pheromone[b][n] = (1 - rho) * Pheromone[b][n] + rho * diffusion
    #         plot.plot([x[n], x[b]], [y[n], y[b]], c='k', alpha=0.3, lw=3, ls='--')
    # plot.plot([x[0], x[0]], [y[0], y[0]], c='k', alpha=0.3, lw=3, ls='--', label='Neighbor trail')
    # plt.legend()
    # plt.tight_layout()
    # plt.savefig('edge_diffuse.pdf', dpi=600)
    # plt.savefig('edge_diffuse.png', dpi=600)
    #
    # plt.show()
    # quit(-1)


def pheromone_global_update(route, scope, delta):
    for idx in range(1, len(route) - 1):
        a, b = route[idx-1], route[idx]
        Pheromone[a][b] = (1 - rho) * Pheromone[a][b] + rho * delta
        Pheromone[b][a] = (1 - rho) * Pheromone[b][a] + rho * delta


def pheromone_vision(plot):
    dot_size = round(np.log2(instance.n)) * 2
    x, y = instance.x, instance.y
    plot.scatter(x[1:], y[1:], s=dot_size, c='grey')
    plot.xticks([])
    plot.yticks([])
    alpha = normalization(Pheromone)
    for i in range(1, instance.n+1):
        for j in range(1, instance.n+1):
            plot.plot([x[i], x[j]], [y[i], y[j]], c='k', alpha=alpha[i][j])
    return plot


class AntTeam:
    def __init__(self):
        self.routes = None
        self.length = None
        self.lengths = None

    def construct(self):
        self.routes = [[D] for _ in range(m)]
        intentions = np.zeros(m, dtype=int)
        cities = deepcopy(V)
        while cities:
            intentions[:] = 0
            for _m in range(m):
                curr = self.routes[_m][-1]
                candidates = set(cities) & set(cities_of(_m))
                intentions[_m] = transition(curr, candidates)
            for i in range(m):
                for j in range(i+1, m):
                    if intentions[i] == intentions[j]:
                        cost_i = Distance[self.routes[i][-1]][intentions[i]]
                        cost_j = Distance[self.routes[j][-1]][intentions[j]]
                        if cost_i > cost_j:
                            intentions[i] = 0
                        else:
                            intentions[j] = 0
            for _m in range(m):
                if intentions[_m] != 0:
                    self.routes[_m].append(intentions[_m])
                    # local pheromone update
                    a, b = self.routes[_m][-1], self.routes[_m][-2]
                    Pheromone[a][b] = (1 - xee) * Pheromone[a][b] + xee * tau_0
                    Pheromone[b][a] = (1 - xee) * Pheromone[b][a] + xee * tau_0
                    cities.remove(intentions[_m])
        for _m in range(m):
            self.routes[_m].append(D)
        self.calc_length()

    def sequential_build(self):
        self.routes = [[D] for _ in range(m)]
        cities = deepcopy(V)
        while cities:
            for _m in range(m):
                curr = self.routes[_m][-1]
                candidates = set(cities) & set(cities_of(_m))
                tar = transition(curr, candidates)
                if tar == 0:
                    continue
                self.routes[_m].append(tar)
                # local pheromone update
                a, b = self.routes[_m][-1], self.routes[_m][-2]
                Pheromone[a][b] = (1 - xee) * Pheromone[a][b] + xee * tau_0
                Pheromone[b][a] = (1 - xee) * Pheromone[b][a] + xee * tau_0
                cities.remove(tar)
        for _m in range(m):
            self.routes[_m].append(D)
        self.calc_length()

    def calc_length(self):
        self.lengths = [0 for _ in range(m)]
        for r, route in enumerate(self.routes):
            for idx in range(1, len(route)):
                self.lengths[r] += Distance[route[idx - 1]][route[idx]]
        self.length = sum(self.lengths)

    def minimal_insert(self, city):
        r, idx, cost = -1, -1, 1e99
        M = salesmen_of(city)
        for _r in M:
            for _idx in range(1, len(self.routes[_r])):
                p, q = self.routes[_r][_idx - 1], self.routes[_r][_idx]
                _cost = Distance[p][city] + Distance[city][q] - Distance[p][q]
                if cost > _cost:
                    r, idx, cost = _r, _idx, _cost
        self.routes[r].insert(idx, city)
        self.lengths[r] += cost
        self.length = sum(self.lengths)

    def minimal_init(self):
        self.routes = [[D, D] for _ in range(m)]
        self.lengths = [0 for _ in range(m)]
        cities = deepcopy(V)
        shuffle(cities)
        for c in cities:
            self.minimal_insert(c)

    def local_search_2_opt(self):
        for _m in range(m):
            route = self.routes[_m]
            L = len(route)
            for i in range(1, L - 1):
                for j in range(i + 1, L - 1):
                    new_route = route[:i] + route[j:i-1:-1] + route[j+1:]
                    new_length = 0
                    for k in range(1, L):
                        new_length += Distance[new_route[k-1]][new_route[k]]
                    if new_length < self.lengths[_m]:
                        self.routes[_m] = new_route
                        self.lengths[_m] = new_length


problem = 'VCMTSP'
method = 'ACO'
feature = 'base'
# plt.figure(figsize=(10, 4))
max_epoch = 1
max_gen = 1000
N = 10  # number of teams
alpha = 1  # pheromone weight
beta = 2  # heuristic weight
rho = 0.1
xee = 0.1
q_zero = 0.9
scope = 1
team = AntTeam()
team.minimal_init()
tau_0 = 1.0 / (instance.n * team.length)
Pheromone = np.full(shape=(instance.n+1, instance.n+1), fill_value=tau_0)
# PheromoneUpdate = pheromone_diffusion
PheromoneUpdate = pheromone_global_update
if __name__ == '__main__':
    date = strftime("%Y-%m-%d", localtime())
    save_path = f'./results/{date}/{case}/{m}'
    make_dir(save_path)
    save_name = f'{method}_{feature}'
    start_time = time()
    optimum = AntTeam()
    optimum.length = 1e99
    for epoch in range(1, max_epoch + 1):
        # Initialization.
        Pheromone = np.full(shape=(instance.n+1, instance.n+1), fill_value=tau_0)
        # Progress
        gBest_list = []
        pBest_list = []
        pWorst_list = []
        gBest = AntTeam()
        gBest.length = 1e99
        for gen in range(1, max_gen + 1):
            colony = []
            for i in range(N):
                team = AntTeam()
                team.construct()
                # team.sequential_build()
                team.local_search_2_opt()
                colony.append(team)
            pBest = deepcopy(min(colony, key=lambda _: _.length))
            if gBest.length > pBest.length:
                gBest = deepcopy(pBest)
            # global pheromone update
            for route in gBest.routes:
                PheromoneUpdate(route, scope=scope, delta=(1.0 / gBest.length))
            # record evo info
            pBest_list.append(pBest.length)
            gBest_list.append(gBest.length)
            pWorst_list.append(max(colony, key=lambda _: _.length).length)
            # report
            print(f'problem: {problem} {case} m={m}, method: {method}, feature: {feature}, beta={beta}, rho={rho}, xee={xee}, q0={q_zero}. '
                  f'At epoch {epoch} gen {gen}: {round(gBest.length, 2)}', end='\r')

            if gen == 500:
                plt.clf()
                plt.imshow(Pheromone)
                plt.colorbar()
                plt.tight_layout()
                plt.savefig('pheromone_matrix_global.pdf', dpi=600)
                plt.show()

            # plot curves & routes
        # plt.clf()
        # plt.subplot(1, 2, 1)
        # # pheromone_vision(plt)
        # plt.imshow(Pheromone)
        # plt.plot(range(len(gBest_list)), gBest_list, label='gBest')
        # # plt.plot(range(len(pBest_list)), pBest_list, label='pBest')
        # # plt.plot(range(len(pWorst_list)), pWorst_list, label='pWorst')
        # plt.legend()
        # plt.title(f'PD total length: {gBest.length:.2f}')
        # plt.subplot(1, 2, 2)
        # instance.draw(plt, routes=gBest.routes)
        # plt.pause(0.001)

        # plt.show()

        if optimum.length > gBest.length:
            optimum = deepcopy(gBest)
        # save info per epoch
        mode = 'w' if epoch == 1 else 'ab'
        with open(rf'{save_path}/{save_name}_fitness.txt', 'w') as f:
            np.savetxt(f, np.array(gBest_list).reshape(1, len(gBest_list)))
        # end report
        print(f'problem: {problem} {case} m={m}, method: {method}, feature: {feature}, beta={beta}, rho={rho}, xee={xee}, q0={q_zero}. '
              f'At epoch {epoch} gen {max_gen}: {round(gBest.length, 2)}')
    # save optimal route ==============================
    route = []
    for _ in optimum.routes:
        route.extend(_)
    route = np.array(route)
    with open(rf'{save_path}/{save_name}_route.txt', 'w') as f:
        np.savetxt(f, route.reshape(1, len(route)), fmt='%d')
    # print time cost
    print(f"Time cost {strftime('%H:%M:%S', gmtime(time() - start_time))}")
